import cv2
import random
import numpy as np

class BrightnessJitterAug():
    def __init__(self, brightness):
        self.brightness = brightness
    def __call__(self, src):
        alpha = 1.0 + random.uniform(-self.brightness, self.brightness)
        src *= alpha
        return src

class ContrastJitterAug():
    def __init__(self, contrast):
        self.contrast = contrast
        self.coef = np.array([[[0.299, 0.587, 0.114]]])
    def __call__(self, src):
        alpha = 1.0 + random.uniform(-self.contrast, self.contrast)
        gray = src * self.coef
        gray = (3.0 * (1.0 - alpha) / np.prod(gray.shape)) * np.sum(gray)
        src *= alpha
        src += gray
        return src

class SaturationJitterAug():
    def __init__(self, saturation):
        self.saturation = saturation
        self.coef = np.array([[[0.299, 0.587, 0.114]]])
    def __call__(self, src):
        alpha = 1.0 + random.uniform(-self.saturation, self.saturation)
        gray = src * self.coef
        gray = np.sum(gray, axis=2, keepdims=True)
        gray *= (1.0 - alpha)
        src *= alpha
        src += gray
        return src

class HueJitterAug():
    def __init__(self, hue):
        self.hue = hue
        self.tyiq = np.array([[0.299, 0.587, 0.114],
                            [0.596, -0.274, -0.321],
                            [0.211, -0.523, 0.311]])
        self.ityiq = np.array([[1.0, 0.956, 0.621],
                            [1.0, -0.272, -0.647],
                            [1.0, -1.107, 1.705]])
    def __call__(self, src):
        alpha = random.uniform(-self.hue, self.hue)
        u = np.cos(alpha * np.pi)
        w = np.sin(alpha * np.pi)
        bt = np.array([[1.0, 0.0, 0.0],
                        [0.0, u, -w],
                        [0.0, w, u]])
        t = np.dot(np.dot(self.ityiq, bt), self.tyiq).T
        src = np.dot(src, np.array(t))
        return src



class ColorJitterAug():
    def __init__(self, brightness, contrast, saturation):
        self.ts = []
        if brightness > 0:
            self.ts.append(BrightnessJitterAug(brightness))
        if contrast > 0:
            self.ts.append(ContrastJitterAug(contrast))
        if saturation > 0:
            self.ts.append(SaturationJitterAug(saturation))

    def __call__(self, src):
        random.shuffle(self.ts)
        for t in self.ts:
            src = t(src)
        return src


class LightingAug():
    def __init__(self, alphastd):
        self.alphastd = alphastd
        self.eigval = np.array([55.46, 4.794, 1.148])
        self.eigvec = np.array([[-0.5675, 0.7192, 0.4009],
                        [-0.5808, -0.0045, -0.8140],
                        [-0.5836, -0.6948, 0.4203]])

    def __call__(self, src):
        alpha = np.random.normal(0, self.alphastd, size=(3,))
        rgb = np.dot(self.eigvec * alpha, self.eigval)
        src += np.array(rgb)
        return src

class RandomGrayAug():
    def __init__(self, p):
        self.p = p
        self.mat = np.array([[0.21, 0.21, 0.21],
                             [0.72, 0.72, 0.72],
                             [0.07, 0.07, 0.07]])

    def __call__(self, src):
        if random.random() < self.p:
            src = np.dot(src, self.mat)
        return src

def RandomJitterAug(src):
    if src.dtype=='uint8':
        src = src.astype('float32')
    aug_list = [
        ColorJitterAug(0.3, 0.3, 0.3),
        LightingAug(0.1),
        RandomGrayAug(0.2)]
    # random.shuffle(aug_list)
    for aug in aug_list:
        src = aug(src)
    return src


if __name__ == "__main__":
    img = cv2.imread('../../test/1.jpg')
    for i in range(20, 30):
        res = RandomJitterAug(img)
        cv2.imwrite('%d.jpg'%i, res)
